/*
 *   Copyright 2021 zzh
 *
 *   Licensed under the Apache License, Version 2.0 (the "License");
 *   you may not use this file except in compliance with the License.
 *   You may obtain a copy of the License at
 *
 *       http://www.apache.org/licenses/LICENSE-2.0
 *
 *   Unless required by applicable law or agreed to in writing, software
 *   distributed under the License is distributed on an "AS IS" BASIS,
 *   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *   See the License for the specific language governing permissions and
 *   limitations under the License.
 *
 */

package com.aduib.boot.common.buffer;

import com.aduib.boot.common.concurrent.ThreadFactoryBuilder;
import com.aduib.boot.common.concurrent.ThreadPoolBuilder;
import com.aduib.boot.common.concurrent.internal.FastThreadLocal;
import com.aduib.boot.common.concurrent.internal.FastThreadLocalThread;
import com.aduib.boot.common.util.SystemUtils;

import java.util.List;
import java.util.concurrent.CopyOnWriteArrayList;
import java.util.concurrent.Executor;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;

/**
 * @description: PooledBufferAllocator
 * @author: zzh
 * @date: 2021/12/2 16:43
 */
public class PooledBufferAllocator {

  static final int DEFAULT_MAX_CAPACITY = Integer.MAX_VALUE;
  static final ScheduledExecutorService scheduledExecutorService =
      ThreadPoolBuilder.create(
          1,
          ThreadFactoryBuilder.createJava("PooledBuffer_CACHE_TRIM_INTERVAL", false),
          new ThreadPoolExecutor.AbortPolicy());
  private static final int DEFAULT_NUM_DIRECT_ARENA;
  // 8192
  private static final int DEFAULT_PAGE_SIZE = 8192;
  private static final int DEFAULT_MAX_ORDER = 11;
  private static final int DEFAULT_NORMAL_CACHE_SIZE = 64;
  private static final int DEFAULT_MAX_CACHED_BUFFER_CAPACITY = 32 * 1024;
  private static final int DEFAULT_CACHE_TRIM_INTERVAL = 8192;
  private static final long DEFAULT_CACHE_TRIM_INTERVAL_MILLIS = 10;
  private static final boolean DEFAULT_USE_CACHE_FOR_ALL_THREADS = true;
  // 内存对齐，一般是4k
  private static final int DEFAULT_DIRECT_MEMORY_CACHE_ALIGNMENT = 0;
  // 4096
  private static final int MIN_PAGE_SIZE = 4096;
  private static final int MAX_CHUNK_SIZE = (int) (((long) Integer.MAX_VALUE + 1) / 2);

  static {
    final int defaultMinNumArena = SystemUtils.availableProcessors() * 2;
    final int defaultChunkSize = DEFAULT_PAGE_SIZE << DEFAULT_MAX_ORDER;
    DEFAULT_NUM_DIRECT_ARENA =
        Math.max(
            0,
            (int)
                Math.min(
                    defaultMinNumArena, SystemUtils.maxDirectMemory() / defaultChunkSize / 2 / 3));
  }

  private final BufferPool[] directArenas;
  private final int normalCacheSize;
  private final PoolThreadLocalCache threadCache;
  private final List<PoolThreadCache> poolThreadCacheList = new CopyOnWriteArrayList<>();
  private final Runnable trimTask = PooledBufferAllocator.this::trimCurrentThreadCache;
  private final int chunkSize;

  public static final PooledBufferAllocator DEFAULT = new PooledBufferAllocator();

  public PooledBufferAllocator() {
    this(DEFAULT_NUM_DIRECT_ARENA, DEFAULT_PAGE_SIZE);
  }

  public PooledBufferAllocator(int nDirectArena, int pageSize) {
    this(
        nDirectArena,
        DEFAULT_PAGE_SIZE,
        DEFAULT_MAX_ORDER,
        DEFAULT_NORMAL_CACHE_SIZE,
        DEFAULT_USE_CACHE_FOR_ALL_THREADS,
        DEFAULT_DIRECT_MEMORY_CACHE_ALIGNMENT);
  }

  public PooledBufferAllocator(
      int nDirectArena,
      int pageSize,
      int maxOrder,
      int normalCacheSize,
      boolean useCacheForAllThreads,
      int directMemoryCacheAlignment) {
    this.threadCache = new PoolThreadLocalCache(useCacheForAllThreads);
    this.normalCacheSize = normalCacheSize;
    chunkSize = validateAndCalculateChunkSize(pageSize, maxOrder);

    int pageShifts = validateAndCalculatePageShifts(pageSize, directMemoryCacheAlignment);

    if (nDirectArena > 0) {
      directArenas = newArenaArray(nDirectArena);
      for (int i = 0; i < directArenas.length; i++) {
        BufferPoolArena arena =
            new BufferPoolArena(chunkSize, pageSize, pageShifts, 0, directMemoryCacheAlignment);
        directArenas[i] = arena;
      }
    } else {
      directArenas = null;
    }
  }

  /**
   * Returns the current {@link Executor} that uses the {@link Thread}, or {@code null} if none /
   * unknown.
   */
  public static ScheduledExecutorService currentExecutor() {
    return scheduledExecutorService;
  }

  @SuppressWarnings("unchecked")
  private static BufferPool[] newArenaArray(int size) {
    return new BufferPool[size];
  }

  private static int validateAndCalculatePageShifts(int pageSize, int alignment) {
    if (pageSize < MIN_PAGE_SIZE) {
      throw new IllegalArgumentException(
          "pageSize: " + pageSize + " (expected: " + MIN_PAGE_SIZE + ')');
    }

    if ((pageSize & pageSize - 1) != 0) {
      throw new IllegalArgumentException("pageSize: " + pageSize + " (expected: power of 2)");
    }

    if (pageSize < alignment) {
      throw new IllegalArgumentException(
          "Alignment cannot be greater than page size. "
              + "Alignment: "
              + alignment
              + ", page size: "
              + pageSize
              + '.');
    }

    // Logarithm base 2. At this point we know that pageSize is a power of two.
    return Integer.SIZE - 1 - Integer.numberOfLeadingZeros(pageSize);
  }

  private static int validateAndCalculateChunkSize(int pageSize, int maxOrder) {
    if (maxOrder > 14) {
      throw new IllegalArgumentException("maxOrder: " + maxOrder + " (expected: 0-14)");
    }

    // Ensure the resulting chunkSize does not overflow.
    int chunkSize = pageSize;
    for (int i = maxOrder; i > 0; i--) {
      if (chunkSize > MAX_CHUNK_SIZE / 2) {
        throw new IllegalArgumentException(
            String.format(
                "pageSize (%d) << maxOrder (%d) must not exceed %d",
                pageSize, maxOrder, MAX_CHUNK_SIZE));
      }
      chunkSize <<= 1;
    }
    return chunkSize;
  }

  /**
   * Default number of direct arenas - System Property: io.netty.allocator.numDirectArenas - default
   * 2 * cores
   */
  public static int defaultNumDirectArena() {
    return DEFAULT_NUM_DIRECT_ARENA;
  }

  /** Default buffer page size - System Property: io.netty.allocator.pageSize - default 8192 */
  public static int defaultPageSize() {
    return DEFAULT_PAGE_SIZE;
  }

  /** Default maximum order - System Property: io.netty.allocator.maxOrder - default 11 */
  public static int defaultMaxOrder() {
    return DEFAULT_MAX_ORDER;
  }

  /**
   * Default thread caching behavior - System Property: io.netty.allocator.useCacheForAllThreads -
   * default true
   */
  public static boolean defaultUseCacheForAllThreads() {
    return DEFAULT_USE_CACHE_FOR_ALL_THREADS;
  }

  /**
   * Default normal cache size - System Property: io.netty.allocator.normalCacheSize - default 64
   */
  public static int defaultNormalCacheSize() {
    return DEFAULT_NORMAL_CACHE_SIZE;
  }

  /** Return {@code true} if direct memory cache alignment is supported, {@code false} otherwise. */
  public static boolean isDirectMemoryCacheAlignmentSupported() {
    return SystemUtils.hasUnsafe();
  }

  private static void validate(int initialCapacity, int maxCapacity) {
    if (initialCapacity > maxCapacity) {
      throw new IllegalArgumentException(
          String.format(
              "initialCapacity: %d (expected: not greater than maxCapacity(%d)",
              initialCapacity, maxCapacity));
    }
  }

  protected Buffer newBuffer(int capacity) {
    PoolThreadCache cache = threadCache.get();
    BufferPool directArena = cache.bufferPool;

    final Buffer buf;
    if (directArena != null) {
      buf = directArena.allocate(capacity);
    } else {
      buf = BufferUtils.allocateBuffer(capacity);
    }

    return buf;
  }

  protected BigBuffer newBigBuffer() {
    PoolThreadCache cache = threadCache.get();
    BufferPool directArena = cache.bufferPool;

    final BigBuffer buf;
    if (directArena != null) {
      buf = directArena.allocateBig();
    } else {
      buf = null;
    }

    return buf;
  }

  public boolean isDirectBufferPooled() {
    return directArenas != null;
  }

  final PoolThreadCache threadCache() {
    PoolThreadCache cache = threadCache.get();
    assert cache != null;
    return cache;
  }

  /**
   * Trim thread local cache for the current {@link Thread}, which will give back any cached memory
   * that was not allocated frequently since the last trim operation.
   *
   * <p>Returns {@code true} if a cache for the current {@link Thread} exists and so was trimmed,
   * false otherwise.
   */
  public boolean trimCurrentThreadCache() {
    boolean flag = false;
    for (PoolThreadCache cache : poolThreadCacheList) {
      if (cache != null) {
        cache.trim();
        flag = false;
      }
      flag = true;
    }
    return flag;
  }

  protected Buffer newPoolBuffer(int initCapacity) {
    validate(initCapacity, DEFAULT_MAX_CAPACITY);

    PoolThreadCache cache = threadCache.get();
    BufferPool directArena = cache.bufferPool;

    final Buffer buf;
    if (directArena != null) {
      buf = directArena.allocate(cache, initCapacity, DEFAULT_MAX_CAPACITY);
    } else {
      buf = BufferUtils.allocateBuffer(initCapacity);
    }

    return buf;
  }

  final class PoolThreadLocalCache extends FastThreadLocal<PoolThreadCache> {
    private final boolean useCacheForAllThreads;

    PoolThreadLocalCache(boolean useCacheForAllThreads) {
      super();
      this.useCacheForAllThreads = useCacheForAllThreads;
    }

    @Override
    protected synchronized PoolThreadCache initialValue() {
      final BufferPool directArena = leastUsedArena(directArenas);

      final Thread current = Thread.currentThread();
      if (useCacheForAllThreads || current instanceof FastThreadLocalThread) {
        final PoolThreadCache cache =
            new PoolThreadCache(
                directArena, normalCacheSize,
                DEFAULT_MAX_CACHED_BUFFER_CAPACITY, DEFAULT_CACHE_TRIM_INTERVAL);

        if (DEFAULT_CACHE_TRIM_INTERVAL_MILLIS > 0) {
          poolThreadCacheList.add(cache);
          final ScheduledExecutorService executor = currentExecutor();
          if (executor != null) {
            executor.scheduleAtFixedRate(
                trimTask,
                DEFAULT_CACHE_TRIM_INTERVAL_MILLIS,
                DEFAULT_CACHE_TRIM_INTERVAL_MILLIS,
                TimeUnit.MILLISECONDS);
          }
        }
        return cache;
      }
      // No caching so just use 0 as sizes.
      return new PoolThreadCache(directArena, 0, 0, 0);
    }

    @Override
    protected void onRemoval(PoolThreadCache threadCache) {
      threadCache.free(false);
    }

    private BufferPool leastUsedArena(BufferPool[] arenas) {
      if (arenas == null || arenas.length == 0) {
        return null;
      }

      BufferPoolArena minArena = (BufferPoolArena) arenas[0];
      for (int i = 1; i < arenas.length; i++) {
        BufferPoolArena arena = (BufferPoolArena) arenas[i];
        if (arena.numThreadCaches.get() < minArena.numThreadCaches.get()) {
          minArena = arena;
        }
      }

      return minArena;
    }
  }
}
